#!/bin/bash
ASCEND_SECLOG="/var/log/ascend_seclog"
pci_dev_path="/sys/bus/pci/devices/"
logFile="${ASCEND_SECLOG}/ascend_install.log"
hotreset_status_file="/var/log/hotreset_status.log"
installInfo="/etc/ascend_install.info"
SHELL_DIR=$(cd "$(dirname "$0")" || exit; pwd)
COMMON_SHELL="$SHELL_DIR/common.sh"
hi1910_bdf="19e5:d100"
hi1910p_bdf="19e5:d500"
hi1980_bdf="19e5:d801"
hi1980b_num=`lspci | grep -c d802`
# load common.sh, get install.info
source "${COMMON_SHELL}"
# read Driver_Install_Path_Param from installInfo
Driver_Install_Path_Param=$(getInstallParam "Driver_Install_Path_Param" "${installInfo}")
Driver_Install_Mode="$(getInstallParam "Driver_Install_Mode" "${installInfo}")"

if [ -f "${Driver_Install_Path_Param}"/driver/script/specific_func.inc ];then
    . "${Driver_Install_Path_Param}"/driver/script/specific_func.inc
elif [ -f  "${PWD}"/driver/script/specific_func.inc ];then
    . ./driver/script/specific_func.inc
fi
log() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    user_id=`id | awk '{printf $1}'`
    echo "[Driver] [$cur_date] [$user_id] "$1 >> $logFile
}
drvEcho() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    echo "[Driver] [$cur_date] $1"
}

# get modules from specific_func.inc
remove_all_ko()
{
    local module modules msg

    # ignore last ko in specific_func.inc
    if [ "${Driver_Install_Mode}" = "normal" ]; then
        modules="$(echo ${mode_normal} | awk -F [\|] '{for(i=1;i<NF;i=i+1) print $i}')"
    elif [ "${Driver_Install_Mode}" = "vnpu_host" ]; then
        modules="$(echo ${mode_vnpu_host} | awk -F [\|] '{for(i=1;i<NF;i=i+1) print $i}')"
    elif [ "${Driver_Install_Mode}" = "vnpu_docker_host" ]; then
        modules="$(echo ${mode_vnpu_docker_host} | awk -F [\|] '{for(i=1;i<NF;i=i+1) print $i}')"
    fi

    for module in $modules
    do
        msg=$(rmmod $module 2>&1)
        if [ $? -ne 0 ] && [ "$module" != "host_notify" ] && [ "$module" != "drv_vascend" ];then
            log "[ERROR]$msg"
            log "[ERROR]rmmod $module failed"
            return 1
        fi
        log "[INFO]rmmod $module success"
    done

    lsmod | grep "drv_seclib_host" > /dev/null 2>&1
    if [ $? -eq 0 ]; then
        log "[INFO]rmmod drv_seclib_host"
        rmmod drv_seclib_host >& /dev/null || { log "[ERROR]rmmod failed" && return 1; }
    fi

    return 0
}
master_set_hotreset_flag()
{
        local bdfs=$1
        local bdf
        for bdf in $bdfs
        do
            log "[INFO]set hot reset flag: "$bdf
            echo 1 > $pci_dev_path$bdf"/hotreset_flag"
        done
}

get_bdfs()
{
    local find_type=$1
    local bdf bdfs

    bdfs=""
    get_pci_info
    for bdf in ${g_bdfs[*]}
    do
        if [ "$find_type" == "all" ];then
            bdfs=$bdfs$bdf" "
        fi
        if [ "$find_type" == "master" ];then
            if [ -e $pci_dev_path$bdf"/chip_id" ];then
                chip_id=`cat $pci_dev_path$bdf"/chip_id"`
                if [ "$chip_id" == "0" ];then
                    bdfs=$bdfs$bdf" "
                fi
            fi
                fi
        if [ "$find_type" == "slave" ];then
                        if [ -e $pci_dev_path$bdf"/chip_id" ];then
                                chip_id=`cat $pci_dev_path$bdf"/chip_id"`
                                if [ "$chip_id" != "0" ];then
                                        bdfs=$bdfs$bdf" "
                                fi
                        fi
                fi
    done

    echo "$bdfs"

    return 0;
}

get_bridge_bus_bdfs()
{
    local bdf bdfs bus_bdfs;
    bdfs="$1"
    bus_bdfs=""
        for bdf in $bdfs
        do
        if [ -e $pci_dev_path$bdf"/bus_name" ];then
            bus_name=`cat $pci_dev_path$bdf"/bus_name"`
            bus_bdfs=$bus_bdfs$bus_name" "
        fi
    done
    echo "$bus_bdfs"
}

remove_devs()
{
    local bdfs=$1
    local bdf=""

    for bdf in ${bdfs}
    do
        if echo ${bdf} | grep -Eq "^0{4}:\w{2}:\w{2}.2$"; then
            log "[INFO]remove dev: ${bdf}"
            echo 1 > ${pci_dev_path}${bdf}"/remove"

            bdf="${bdf/.2/.1}"
            log "[INFO]remove dev: ${bdf}"
            echo 1 > ${pci_dev_path}${bdf}"/remove"

            bdf=${bdf/.1/.0}
            log "[INFO]remove dev: ${bdf}"
            echo 1 > ${pci_dev_path}${bdf}"/remove"
        else
            log "[INFO]remove dev: ${bdf}"
            echo 1 > ${pci_dev_path}${bdf}"/remove"
        fi
    done

    return 0
}

check_exist_reset_file()
{
    local bdf bus_bdfs
    bus_bdfs="$1"

    for bdf in $bus_bdfs
    do
        if [ ! -e ${pci_dev_path}${bdf}"/reset" ]; then
            log "[WARNING]${pci_dev_path}${bdf}/reset is not exist"
            return 1
        fi
    done

    return 0
}

master_bridge_do_hotreset()
{
    local bdf bus_bdfs
    bus_bdfs="$1"

    for bdf in $bus_bdfs
        do
            log "[INFO]set hot_reset bus(43): "$bdf
            setpci -s "$bdf" 3e.b=43
        done
    sleep 0.002
    for bdf in $bus_bdfs
        do
            log "[INFO]set hot_reset bus(03): "$bdf
            setpci -s "$bdf" 3e.b=03
        done
}
master_do_hotreset()
{
    local bdfs=$1
    local bdf=""

    for bdf in ${bdfs}
    do
        if echo ${bdf} | grep -Eq "^0{4}:\w{2}:\w{2}.2$"; then
            master_bridge_do_hotreset "$masterDevsBus"
            sleep 1
            log "[INFO].remove master device to redistribution bar address"
            remove_devs "$masterDevs"
            break
        else
            log "[INFO]hot reset dev: "${bdf}
            echo 1 > ${pci_dev_path}${bdf}"/reset"
        fi
    done
}
rescan_bus()
{
    local bus_bdfs=$1
    local bdf
    for bdf in $bus_bdfs
    do
        log "[INFO]rescan bus: "$bdf
        if [ -e $pci_dev_path$bdf"/rescan" ]; then
            echo 1 > $pci_dev_path$bdf"/rescan"
        else
            echo 1 > $pci_dev_path$bdf"/dev_rescan"
        fi
    done
}

get_devs_num()
{
    local bdf bdfs
    bdfs=$1
    i=0
        for bdf in $bdfs
        do
        ((i++))
        done
    echo $i
}

check_and_disable_sriov()
{
    local bus_bdfs=$1
    local bdf
    for bdf in $bus_bdfs
    do
        log "[INFO]check sriov support status: $bdf"
        lspci -vvvs $bdf | grep SR-IOV
        if [ $? -eq 0 ]; then
            vf_num=`cat $pci_dev_path$bdf/sriov_numvfs`
            if [ $vf_num != "0" ];then
                log "[INFO]disable sriov for "$bdf
                echo 0 > $pci_dev_path$bdf"/sriov_numvfs"
                if [ $? -ne 0 ]; then
                    log "[ERROR] disable sriov failed"
                    return 1
               fi
            else
                log "[INFO] sriov is disable："$bdf
            fi
        else
            log "[INFO]sriov is not support for "$bdf
        fi
    done
}

hot_reset()
{
    if [ -e "$hotreset_status_file" ];then
        hotreset_status=`cat "$hotreset_status_file"`
        hotreset_status=${hotreset_status%.*}
    else
        hotreset_status="unknown"
    fi
     log "[INFO]hotreset_status:"$hotreset_status
    if [ "$hotreset_status" = "scan_success" ];then
        log "[INFO]hot reset has already executed"
        return
    fi

    log "[INFO]try hot reset device"
    echo "start" > "$hotreset_status_file"

    checkUserProcess >& /dev/null
        res=`echo $?`   # get return result
        if [ $res = "1" ]; then
                log "[INFO]user process exist, no hot reset"
                echo "abort" > "$hotreset_status_file"
                return
        fi
    if [ -d "/sys/bus/mdev/devices/" ] && [ "$(ls -A /sys/bus/mdev/devices/)" ]; then
        log "[INFO]user process exist, no hot reset"
        echo "abort" > "$hotreset_status_file"
        return
    fi
    local allDevs=`get_bdfs "all"`
    log "[INFO]allDevs: $allDevs"
    local masterDevs=`get_bdfs "master"`
    log "[INFO]masterDevs: $masterDevs"
    local slaveDevs=`get_bdfs "slave"`
    log "[INFO]slaveDevs: $slaveDevs"

    local devnum=`get_devs_num "$allDevs"`
    if [ "$pcie_bdf"x = "$hi1980_bdf"x ] || [ $hi1980b_num -ne 0 ]; then
        log "[INFO]devnum: $devnum"
        local masterDevsBus=`get_bridge_bus_bdfs "$masterDevs"`
        log "[INFO]masterDevsBus: $masterDevsBus"
        local slaveDevsBus=`get_bridge_bus_bdfs "$slaveDevs"`
        log "[INFO]slaveDevsBus: $slaveDevsBus"

        local busnum=`get_devs_num "$masterDevsBus"`
        if [ $busnum == 0 ]; then
            log "[INFO]master bus num: "$busnum" not surport hotreset"
            echo "abort" > "$hotreset_status_file"
            return
        fi
    fi

    check_exist_reset_file "$allDevs"
    res=`echo $?`   # get return result
    if [ $res = "1" ]; then
        echo "fail" > "$hotreset_status_file"
        return
    fi

    drvEcho "[INFO]device hot reset start"

    log "[INFO]1.set hotreset flag for bios"
    master_set_hotreset_flag "$masterDevs"

    log "[INFO]1.1 check sriov and disable sriov when enable"
    check_and_disable_sriov "$masterDevs" > /dev/null 2>&1
    if [ $? -ne 0 ]; then
        log "[ERROR] check and disable sriov failed"
        echo "abort" > "$hotreset_status_file"
        return
    fi

    local allDevs=`get_bdfs "all"`
    local masterDevs=`get_bdfs "master"`
    local slaveDevs=`get_bdfs "slave"`
    local devnum=`get_devs_num "$allDevs"`

    log "[INFO]2.remove old version modules ko"
    remove_all_ko
    res=`echo $?`   # get return result
    if [ $res = "1" ]; then
        log "[WARNING]ko can not be removed"
        echo "ko_abort" > "$hotreset_status_file"
        return
    fi

    if [ "$pcie_bdf"x = "$hi1980_bdf"x ] || [ $hi1980b_num -ne 0 ]; then
        log "[INFO]3.remove slave devices avoid host reset"
        remove_devs "$slaveDevs"
        sleep 1

        log "[INFO]4.devices do hot reset"
        master_do_hotreset "$masterDevs"
        sleep 2
    else
        log "[INFO]4.devices do hot reset"
        master_do_hotreset "$allDevs"
        sleep 2
    fi

    if [ $hi1980b_num -ne 0 ]; then
        log "[INFO]5.remove master devices prepare for rescan"
        remove_devs "$masterDevs"
        sleep 1
    fi

    log "[INFO]6.rescan and wait all removed devices"
    local scan_devnum try_cnt
    scan_devnum=0
    try_cnt=0
    while [ $devnum != $scan_devnum ]
    do
        if (($try_cnt > 10)); then
            log "[WARNING]hot reset failed, all device num:"$devnum", rescan device num:"$scan_devnum
            echo "fail" > "$hotreset_status_file"
            return
        fi
        ((try_cnt++))

        sleep 1
        if [ "$pcie_bdf"x = "$hi1980_bdf"x ] || [ $hi1980b_num -ne 0 ]; then
            rescan_bus "$masterDevsBus"
            rescan_bus "$slaveDevsBus"
        fi
        allDevs=`get_bdfs "all"`
        scan_devnum=`get_devs_num "$allDevs"`

    done

    drvEcho "[INFO]device hot reset finish"
    log "[INFO]device hot reset finish"
    echo "scan_success."$(date +%s) > "$hotreset_status_file"
}

hot_reset

